"""Shared pipeline execution context."""

from __future__ import annotations

from typing import Optional

import numpy as np

from ..qt_compat import cv2


class PipelineContext:
    """Holds the working image and related frequency domain state."""

    def __init__(self, image: np.ndarray, source_bgr: Optional[np.ndarray] = None) -> None:
        self.image = image.astype(np.float32)
        self.source_bgr: Optional[np.ndarray] = source_bgr
        self.fourier_shift: Optional[np.ndarray] = None
        self.fourier_source: Optional[np.ndarray] = None

    def ensure_frequency(self) -> np.ndarray:
        if self.fourier_shift is None:
            freq_input = self.image
            dft = cv2.dft(freq_input, flags=cv2.DFT_COMPLEX_OUTPUT)
            self.fourier_shift = np.fft.fftshift(dft)
            self.fourier_source = freq_input.copy()
        return self.fourier_shift

    def update_from_frequency(self) -> np.ndarray:
        if self.fourier_shift is None:
            self.ensure_frequency()
        ishift = np.fft.ifftshift(self.fourier_shift)
        img_back = cv2.idft(ishift)
        magnitude = cv2.magnitude(img_back[:, :, 0], img_back[:, :, 1])
        self.image = magnitude
        return magnitude
